import torch
from algorithms.fc3d_sequence_lstm import FC3DSequenceLSTM

def load_compatible_3d_sequence_model(model_path, device=None):
    """
    加载兼容版3D序列LSTM模型
    
    Args:
        model_path (str): 模型文件路径
        device (torch.device, optional): 设备
        
    Returns:
        FC3DSequenceLSTM: 加载的模型
    """
    if device is None:
        device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    
    # 创建模型实例
    model = FC3DSequenceLSTM(
        input_dim=6,      # 3个数字特征 + 3个区域转换特征
        hidden_dim=128,
        num_layers=3,
        dropout=0.3
    )
    
    # 加载模型状态
    checkpoint = torch.load(model_path, map_location=device)
    
    # 检查检查点格式
    if isinstance(checkpoint, dict):
        if 'model_state_dict' in checkpoint:
            model.load_state_dict(checkpoint['model_state_dict'])
        elif 'state_dict' in checkpoint:
            model.load_state_dict(checkpoint['state_dict'])
        elif 'model' in checkpoint:
            model.load_state_dict(checkpoint['model'])
        else:
            # 假设检查点直接包含状态字典
            model.load_state_dict(checkpoint)
    else:
        # 直接加载状态字典
        model.load_state_dict(checkpoint)
    
    # 将模型移动到指定设备
    model = model.to(device)
    model.eval()
    
    return model

if __name__ == "__main__":
    print("[INFO] 兼容版3D序列LSTM模型加载器")
    print("[INFO] 请使用 load_compatible_3d_sequence_model 函数加载模型")